from cryptography import x509
from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
import datetime, os, ipaddress
from typing import cast
from config_types import ServerCertConfig


def generate_server_cert(config: ServerCertConfig, root_path: str = ""):
    """
    生成服务器证书
    
    Args:
        config: 配置字典，包含以下键：
            - inter_key_path: 中间 CA 私钥路径
            - inter_crt_path: 中间 CA 证书路径
            - inter_password: 中间 CA 私钥密码
            - root_crt_path: 根 CA 证书路径
            - output_dir: 输出目录
            - key_path: 私钥文件路径
            - crt_path: 证书文件路径
            - fullchain: 完整证书链文件路径
            - country_name: 国家名称
            - state_or_province_name: 省/州名称
            - locality_name: 城市名称
            - organization_name: 组织名称
            - common_name: 通用名称
            - dns_names: DNS 名称列表
            - ip_addrs: IP 地址列表
            - days: 有效期（天数）
            - key_size: RSA密钥长度（可选，默认2048位）
        root_path: 根路径，会与配置中的路径拼接
    """
    INTER_KEY_PATH = os.path.join(root_path, config["inter_key_path"]) if root_path else config["inter_key_path"]
    INTER_CRT_PATH = os.path.join(root_path, config["inter_crt_path"]) if root_path else config["inter_crt_path"]
    INTER_PASS = config["inter_password"].encode()
    ROOT_CRT_PATH = os.path.join(root_path, config["root_crt_path"]) if root_path else config["root_crt_path"]

    KEY_PATH = os.path.join(root_path, config["key_path"]) if root_path else config["key_path"]
    CRT_PATH = os.path.join(root_path, config["crt_path"]) if root_path else config["crt_path"]
    FULLCHAIN = os.path.join(root_path, config["fullchain"]) if root_path else config["fullchain"]

    COUNTRY_NAME = config["country_name"]
    STATE_OR_PROVINCE_NAME = config["state_or_province_name"]
    LOCALITY_NAME = config["locality_name"]
    ORGANIZATION_NAME = config["organization_name"]
    COMMON_NAME = config["common_name"]

    DNS_NAMES = config["dns_names"]
    IP_ADDRS = config["ip_addrs"]
    DAYS = config["days"]
    KEY_SIZE = config.get("key_size", 2048)  # 默认2048位

    # 确保输出目录存在
    if os.path.dirname(KEY_PATH):
        os.makedirs(os.path.dirname(KEY_PATH), exist_ok=True)
    if os.path.dirname(CRT_PATH):
        os.makedirs(os.path.dirname(CRT_PATH), exist_ok=True)
    if os.path.dirname(FULLCHAIN):
        os.makedirs(os.path.dirname(FULLCHAIN), exist_ok=True)

    # 检查必需的证书文件是否存在
    if not os.path.exists(INTER_KEY_PATH):
        raise FileNotFoundError(f"中间 CA 私钥文件不存在: {INTER_KEY_PATH}")
    if not os.path.exists(INTER_CRT_PATH):
        raise FileNotFoundError(f"中间 CA 证书文件不存在: {INTER_CRT_PATH}")
    if not os.path.exists(ROOT_CRT_PATH):
        raise FileNotFoundError(f"根 CA 证书文件不存在: {ROOT_CRT_PATH}")

    inter_key = serialization.load_pem_private_key(
        open(INTER_KEY_PATH, "rb").read(), password=INTER_PASS
    )
    inter_cert = x509.load_pem_x509_certificate(open(INTER_CRT_PATH, "rb").read())
    
    # 类型断言：我们知道这是 RSA 私钥（因为 make_intermediate_ca.py 生成的是 RSA 密钥）
    inter_key = cast(rsa.RSAPrivateKey, inter_key)

    key = rsa.generate_private_key(public_exponent=65537, key_size=KEY_SIZE)
    subject = x509.Name(
        [
            x509.NameAttribute(NameOID.COUNTRY_NAME, COUNTRY_NAME),
            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, STATE_OR_PROVINCE_NAME),
            x509.NameAttribute(NameOID.LOCALITY_NAME, LOCALITY_NAME),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, ORGANIZATION_NAME),
            x509.NameAttribute(NameOID.COMMON_NAME, COMMON_NAME),
        ]
    )

    sans = [x509.DNSName(d) for d in DNS_NAMES] + [
        x509.IPAddress(ipaddress.ip_address(i)) for i in IP_ADDRS
    ]
    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(inter_cert.subject)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.utcnow())
        .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=DAYS))
        # ---- 基础扩展 ----
        .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        # ---- Key Usage：用于 TLS 握手的签名与加密 ----
        .add_extension(
            x509.KeyUsage(
                digital_signature=True,  # 允许签名（握手签名）
                content_commitment=False,
                key_encipherment=True,  # 允许加密（密钥交换）
                data_encipherment=False,
                key_agreement=False,
                key_cert_sign=False,
                crl_sign=False,
                encipher_only=False,
                decipher_only=False,
            ),
            critical=True,
        )
        # ---- Extended Key Usage：限制用途 ----
        .add_extension(
            x509.ExtendedKeyUsage(
                [
                    ExtendedKeyUsageOID.SERVER_AUTH,
                    ExtendedKeyUsageOID.CLIENT_AUTH,
                ]
            ),
            critical=False,
        )
        # ---- SAN：多域名/IP ----
        .add_extension(x509.SubjectAlternativeName(sans), critical=False)
        # ---- 标识符扩展 ----
        # SubjectKeyIdentifier: 本证书的公钥标识符
        .add_extension(
            x509.SubjectKeyIdentifier.from_public_key(key.public_key()),
            critical=False,
        )
        # AuthorityKeyIdentifier: 标识签发本证书的中间CA的公钥
        # 从中间CA的公钥计算哈希，验证时用于匹配fullchain中中间CA证书的SubjectKeyIdentifier
        # ⚠️ 在TLS握手时，服务器发送的fullchain中，此AKI必须与中间CA证书的SKI匹配
        #    因此，如果重新生成中间CA（新密钥），要么继续使用旧中间CA证书，要么重新签发Server证书
        .add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(
                cast(rsa.RSAPublicKey, inter_cert.public_key())
            ),
            critical=False,
        )
        .sign(inter_key, hashes.SHA256())
    )

    with open(KEY_PATH, "wb") as f:
        f.write(
            key.private_bytes(
                serialization.Encoding.PEM,
                serialization.PrivateFormat.TraditionalOpenSSL,
                serialization.NoEncryption(),
            )
        )
    with open(CRT_PATH, "wb") as f:
        f.write(cert.public_bytes(serialization.Encoding.PEM))

    # 读取证书文件并生成完整链
    root_crt = open(ROOT_CRT_PATH, "rb").read()
    inter_crt = open(INTER_CRT_PATH, "rb").read()
    server_crt = open(CRT_PATH, "rb").read()
    
    # 确保 fullchain 目录存在
    if os.path.dirname(FULLCHAIN):
        os.makedirs(os.path.dirname(FULLCHAIN), exist_ok=True)
    
    with open(FULLCHAIN, "wb") as f:
        f.write(server_crt + inter_crt + root_crt)

    print("Server 证书生成成功：", CRT_PATH)
    print("🔗 Fullchain: ", FULLCHAIN)
    return KEY_PATH, CRT_PATH, FULLCHAIN


if __name__ == "__main__":
    # 独立运行时从配置文件加载
    import json
    from config_types import ConfigDict
    config_path = os.path.join(os.path.dirname(__file__), "./configs/test.json")
    with open(config_path, "r", encoding="utf-8") as f:
        full_config = cast(ConfigDict, json.load(f))
        if "server_cert" not in full_config or not full_config["server_cert"]:
            raise ValueError("配置文件中缺少 'server_cert' 字段或列表为空")
        config = full_config["server_cert"][0]
        root_path = full_config.get("root_path", "")
    generate_server_cert(config, root_path)
